package com.it.topN;

import com.it.operator.utils.SourceUtils;
import com.it.pojo.Event;
import com.it.pojo.UrlVisitBean;
import com.it.windows.WindowAggAndProcessTest;
import org.apache.flink.api.common.eventtime.SerializableTimestampAssigner;
import org.apache.flink.api.common.eventtime.WatermarkStrategy;
import org.apache.flink.api.common.state.ListState;
import org.apache.flink.api.common.state.ListStateDescriptor;
import org.apache.flink.api.common.typeinfo.Types;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.streaming.api.datastream.SingleOutputStreamOperator;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.functions.KeyedProcessFunction;
import org.apache.flink.streaming.api.windowing.assigners.SlidingEventTimeWindows;
import org.apache.flink.streaming.api.windowing.time.Time;
import org.apache.flink.util.Collector;

import java.sql.Timestamp;
import java.time.Duration;
import java.util.ArrayList;
import java.util.Comparator;

/**
 * 如果我们使用增量聚合的方式得到url链接的访问量，但是后续的排序输出Top N比较难实现
 *
 * @author code1997
 */
public class TopNTest1 {
    public static void main(String[] args) throws Exception {
        StreamExecutionEnvironment executionEnvironment = StreamExecutionEnvironment.getExecutionEnvironment();
        SingleOutputStreamOperator<Event> sourceStream = SourceUtils.getEventSource(executionEnvironment)
                .assignTimestampsAndWatermarks(WatermarkStrategy.<Event>forBoundedOutOfOrderness(Duration.ZERO).withTimestampAssigner(new SerializableTimestampAssigner<Event>() {
            @Override
            public long extractTimestamp(Event element, long recordTimestamp) {
                return element.timestamp;
            }
        }));
        SingleOutputStreamOperator<UrlVisitBean> urlCountStream = sourceStream.keyBy(event -> event.url)
                .window(SlidingEventTimeWindows.of(Time.seconds(10), Time.seconds(5)))
                .aggregate(new WindowAggAndProcessTest.MyAggFunction(), new WindowAggAndProcessTest.MyProcessFunction());
        urlCountStream.print("url count");

        urlCountStream.keyBy(data -> data.windowEnd)
                .process(new TopNProcessResult(2))
                .print("top n");
        executionEnvironment.execute();
    }

    public static class TopNProcessResult extends KeyedProcessFunction<Long, UrlVisitBean, String> {

        private Integer n;

        //定义状态列表，而且这个对象需要交给flink进行管理，所以不能直接初始化，应该在生命周期函数中进行初始化。
        private ListState<UrlVisitBean> urlViewCountListState;

        public TopNProcessResult(Integer n) {
            this.n = n;
        }

        @Override
        public void open(Configuration parameters) throws Exception {
            urlViewCountListState = getRuntimeContext().getListState(
                    new ListStateDescriptor<UrlVisitBean>("url_count_list", Types.POJO(UrlVisitBean.class))
            );
        }

        @Override
        public void processElement(UrlVisitBean value, KeyedProcessFunction<Long, UrlVisitBean, String>.Context ctx, Collector<String> out) throws Exception {
            urlViewCountListState.add(value);
            //注册一个windowEnd+1ms的定时器
            ctx.timerService().registerEventTimeTimer(ctx.getCurrentKey() + 1);
        }

        @Override
        public void onTimer(long timestamp, KeyedProcessFunction<Long, UrlVisitBean, String>.OnTimerContext ctx, Collector<String> out) throws Exception {
            //定时器处理
            ArrayList<UrlVisitBean> list = new ArrayList<>();
            for (UrlVisitBean urlVisitBean : urlViewCountListState.get()) {
                list.add(urlVisitBean);
            }
            list.sort(new Comparator<UrlVisitBean>() {
                @Override
                public int compare(UrlVisitBean o1, UrlVisitBean o2) {
                    return (int) (o2.count - o1.count);
                }
            });

            StringBuilder result = new StringBuilder();
            Timestamp end = new Timestamp(ctx.getCurrentKey());
            result.append("------------ window [ 窗口结束时间：" + end + "]Top 2--------------------\n");
            for (int i = 0; i < list.size() && i < 2; i++) {
                UrlVisitBean urlInfo = list.get(i);
                result.append("No" + i + ". " + urlInfo.url + " = " + urlInfo.count + "\n");
            }
            result.append("----------------------------------------\n");
            out.collect(result.toString());
        }
    }

}
